# -*- coding: utf-8 -*-            
# @Time    : 2025/5/19 8:44
# @FileName: yk_pj_02.py
# @Software: PyCharm

import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

# 项目2：RNN股票预测
# 一、项目背景
# 本项目旨在开发一个基于深度学习RNN（循环神经网络）模型的股票预测系统，
# 能够处理历史股票数据并进行未来股票价格的预测。
# 项目的主要步骤包括数据处理、模型构建、模型训练和结果可视化。
# 二、功能需求
# 1.
# 数据预处理
# 1)    数据必须从CSV文件data - 02 - stock_daily.csv中加载，且该文件应以逗号作为分隔符。
# 2)    数据需要被反转（data[::-1]），可能是为了使得时间序列数据按照时间顺序排列
#       （如果原始数据是逆序存储的话）。
# 3)    数据需要经过MinMaxScaler进行归一化处理。
# 4)    序列长度（c）被设定为7，意味着每个输入序列包含7个连续的数据点，而目标值（y_data）是序列之后的那个数据点的最后一个特征值。
data = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
data = data[::-1]
data = MinMaxScaler().fit_transform(data)

c = 7
x = []
y = []
for i in range(len(data) - c):
    x.append(data[i:i + c])
    y.append(data[i + c][-1])

x = torch.tensor(x, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.float32).reshape(-1,1)

# 2.
# 数据划分
# 1)    数据集需要被划分为训练集和测试集，且划分时不能进行随机打乱（shuffle = False）。
x_train, x_test, y_train, y_test = train_test_split(x, y,
                                                    shuffle=False)
input_size = x_train.shape[2]



# 3.
# 模型结构
# 1)    模型必须是一个循环神经网络（RNN），包含两层RNN层和一个全连接层。
class RNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 2)    第一层RNN的input_size必须与输入数据的特征数量相匹配，这里由x.shape[2]决定。
        self.rnn1 = torch.nn.RNN(input_size=input_size, hidden_size=5)
        # 3)    第一层RNN的hidden_size为5，第二层RNN的hidden_size为10。
        self.rnn2 = torch.nn.RNN(input_size=5, hidden_size=10)
        # 4)    全连接层的输入特征数为10，输出特征数为1，用于预测单个值。
        self.fc = torch.nn.Linear(in_features=10, out_features=1)

    def forward(self,x):
        x,y = self.rnn1(x)
        x,y = self.rnn2(x)
        x = self.fc(x[:,-1,:])
        return x

model = RNN()
# # 4.
# # 模型训练
# # 1)    使用均方误差（MSE）作为损失函数。
# # 2)    使用Adam优化器进行参数优化，学习率设定为0.01。
# # 3)    训练过程中，模型需要进行1000次迭代。
# # 4)    在每次迭代中，模型需要进行前向传播、损失计算、反向传播和参数更新。
loss_fn = torch.nn.MSELoss()
op = torch.optim.Adam(model.parameters(), lr=0.01)
loss_list = []
model.train()
for epoch in range(1000):
    op.zero_grad()
    h = model(x_train)
    loss = loss_fn(h, y_train)
    loss_list.append(loss)
    loss.backward()
    op.step()
    if epoch % 10 == 0:
        print(f'epoch:{epoch},loss:{loss}')
# 5.
# 预测与可视化
# 1)    训练完成后，模型需要对测试集进行预测。
# 2)    预测结果需要被重塑为一维数组。
# 3)    使用matplotlib库将预测结果（红色）和实际值（绿色）绘制在同一张图上以便可视化对比。
model.eval()
with torch.no_grad():
    h = model(x_test)
    print(h)
    plt.plot(h, c='r')
    plt.plot(y_test, c='g')
    plt.show()
